# https://github.com/espnet/espnet/tree/master/egs2/aishell/asr1
# https://huggingface.co/pyf98/aishell_ctc_conformer_e15_linear1024

import soundfile
from espnet2.bin.asr_inference import Speech2Text

model_name = "/home/gyf/pkg/conformer_new/aishell_ctc_conformer_e15_linear1024"
wav_file = "/home/gyf/pkg/conformer_new/asr_conformer/speech_recognition/wav/BAC009S0764W0121.wav"

speech2text_kwargs = dict(
        asr_train_config="/home/gyf/pkg/conformer_new/aishell_ctc_conformer_e15_linear1024/exp/asr_train_asr_ctc_conformer_e15_linear1024_raw_zh_char_sp/config.yaml",
        asr_model_file="/home/gyf/pkg/conformer_new/aishell_ctc_conformer_e15_linear1024/exp/asr_train_asr_ctc_conformer_e15_linear1024_raw_zh_char_sp/valid.cer_ctc.ave_10best.pth",
        # transducer_conf=None,
        # lm_train_config=None,
        # lm_file=None,
        # ngram_file=None,
        # token_type=None,
        # bpemodel=None,
        device="cuda",
        maxlenratio=0.0,
        minlenratio=0.0,
        dtype="float32",
        beam_size=20,
        ctc_weight=1.0,
        lm_weight=0.0,
        ngram_weight=0.0,
        penalty=0.0,
        nbest=1,
        # streaming=False,
        # enh_s2t_task=False,
        # multi_asr=False,
        # quantize_asr_model=False,
        # quantize_lm=False,
        # quantize_modules=["Linear"],
        # quantize_dtype="qint8",
        # hugging_face_decoder=False,
        # hugging_face_decoder_max_length=256,
        # time_sync=False,
)

model_tag = None

speech2text = Speech2Text.from_pretrained(
    # model_name,
    # # Decoding parameters are not included in the model file
    # maxlenratio=0.0,
    # minlenratio=0.0,
    # beam_size=20,
    # ctc_weight=0.3,
    # lm_weight=0.5,
    # penalty=0.0,
    # nbest=1
    model_tag = None,
    **speech2text_kwargs,
)


# Confirm the sampling rate is equal to that of the training corpus.
# If not, you need to resample the audio data before inputting to speech2text
speech, rate = soundfile.read(wav_file)
nbests = speech2text(speech)

text, *_ = nbests[0]
print(text)